"""
    test_costfunctions.jl

Test suite for the CostFunctions module of GSICoreAnalysis.jl.
This module tests cost function evaluation, gradient computation,
background and observation terms, and hybrid ensemble methods.
"""

using Test
using GSICoreAnalysis
using GSICoreAnalysis.CostFunctions
using LinearAlgebra
using Random

# Set reproducible seed
Random.seed!(456)

@testset "CostFunctions Module Tests" begin
    
    @testset "CostFunction Construction" begin
        # Test basic construction
        config = AnalysisConfig{Float64}(grid_size = (8, 4, 2))
        cf = CostFunction(config)
        
        @test cf isa CostFunction{Float64}
        @test cf.config === config
        @test size(cf.background_covariance, 1) > 0
        @test size(cf.observation_covariance, 1) > 0
        @test cf.current_cost == 0.0
        @test cf.current_gradient_norm == 0.0
        
        # Test precision specification
        config32 = AnalysisConfig{Float32}(grid_size = (4, 4, 2))
        cf32 = CostFunction{Float32}(config32)
        @test cf32 isa CostFunction{Float32}
        @test eltype(cf32.background_state) == Float32
        
        # Test hybrid cost function
        config_hybrid = AnalysisConfig{Float64}(
            grid_size = (6, 6, 2),
            ensemble_size = 10,
            use_hybrid = true
        )
        hcf = HybridCostFunction(config_hybrid)
        @test hcf isa HybridCostFunction{Float64}
        @test size(hcf.ensemble_perturbations, 2) == 10
        @test hcf.hybrid_coefficient ≈ config_hybrid.hybrid_coeff
    end
    
    @testset "Background Covariance Initialization" begin
        config = AnalysisConfig{Float64}(grid_size = (4, 4, 2))
        cf = CostFunction(config)
        
        n_state = prod(config.grid_size)
        B_test = Diagonal(rand(n_state))
        
        initialize_background_covariance!(cf, B_test)
        @test cf.background_covariance === B_test
        
        # Test hybrid cost function
        hcf = HybridCostFunction(config)
        B_static = Diagonal(rand(n_state))
        initialize_background_covariance!(hcf, B_static)
        @test hcf.static_covariance === B_static
    end
    
    @testset "Ensemble Initialization" begin
        config = AnalysisConfig{Float64}(
            grid_size = (6, 4, 2),
            ensemble_size = 15
        )
        hcf = HybridCostFunction(config)
        
        n_state = prod(config.grid_size)
        perturbations = randn(n_state, config.ensemble_size)
        
        initialize_ensemble!(hcf, perturbations)
        @test hcf.ensemble_perturbations === perturbations
        @test size(hcf.ensemble_perturbations) == (n_state, config.ensemble_size)
    end
    
    @testset "Background and Observations Setup" begin
        config = AnalysisConfig{Float64}(grid_size = (4, 4, 2))
        cf = CostFunction(config)
        
        n_state = prod(config.grid_size)
        background = randn(n_state)
        observations = randn(10)
        
        # Simple observation operator (identity for first few elements)
        H(x) = x[1:length(observations)]
        
        set_background!(cf, background)
        @test cf.background_state ≈ background
        
        set_observations!(cf, observations, H)
        @test cf.observations ≈ observations
        @test cf.observation_operator === H
    end
    
    @testset "Background Term Evaluation" begin
        config = AnalysisConfig{Float64}(grid_size = (4, 4, 2))
        cf = CostFunction(config)
        
        n_state = prod(config.grid_size)
        
        # Set up simple diagonal background covariance
        B_diag = ones(n_state)
        B = Diagonal(B_diag)
        initialize_background_covariance!(cf, B)
        
        # Set background state
        background = zeros(n_state)
        set_background!(cf, background)
        
        # Test with state at background (should give zero cost)
        state = copy(background)
        bg_cost = evaluate_background_term(cf, state)
        @test bg_cost ≈ 0.0 atol=1e-12
        
        # Test with departure from background
        departure = ones(n_state)
        state = background + departure
        bg_cost = evaluate_background_term(cf, state)
        
        # For diagonal B with unit elements: J_b = 0.5 * sum(departure^2) = 0.5 * n_state
        expected_cost = 0.5 * n_state
        @test bg_cost ≈ expected_cost atol=1e-10
        
        # Test scaling property
        scale_factor = 2.0
        scaled_state = background + scale_factor * departure
        scaled_cost = evaluate_background_term(cf, scaled_state)
        @test scaled_cost ≈ scale_factor^2 * bg_cost atol=1e-10
    end
    
    @testset "Hybrid Background Term Evaluation" begin
        config = AnalysisConfig{Float64}(
            grid_size = (4, 4, 2),
            ensemble_size = 8,
            hybrid_coeff = 0.6
        )
        hcf = HybridCostFunction(config)
        
        n_state = prod(config.grid_size)
        
        # Set up static covariance
        B_static = Diagonal(ones(n_state))
        initialize_background_covariance!(hcf, B_static)
        
        # Set up ensemble perturbations
        perturbations = randn(n_state, config.ensemble_size)
        initialize_ensemble!(hcf, perturbations)
        
        # Set background state
        background = zeros(n_state)
        set_background!(hcf, background)
        
        # Test with state at background
        state = copy(background)
        bg_cost = evaluate_background_term(hcf, state)
        @test bg_cost ≈ 0.0 atol=1e-10
        
        # Test with departure from background
        departure = ones(n_state) * 0.1  # Small departure for numerical stability
        state = background + departure
        bg_cost = evaluate_background_term(hcf, state)
        @test bg_cost >= 0.0  # Cost should be non-negative
        @test isfinite(bg_cost)
    end
    
    @testset "Observation Term Evaluation" begin
        config = AnalysisConfig{Float64}(grid_size = (4, 4, 2))
        cf = CostFunction(config)
        
        n_state = prod(config.grid_size)
        n_obs = 5
        
        # Simple observation operator
        H(x) = x[1:n_obs]
        
        # Set up observations and observation error covariance
        observations = randn(n_obs)
        R = Diagonal(ones(n_obs))  # Unit observation error
        
        set_observations!(cf, observations, H)
        cf.observation_covariance = R
        
        # Test with state that matches observations perfectly
        state = zeros(n_state)
        state[1:n_obs] = observations
        obs_cost = evaluate_observation_term(cf, state)
        @test obs_cost ≈ 0.0 atol=1e-12
        
        # Test with state that has unit departure from observations
        state[1:n_obs] = observations .+ 1.0
        obs_cost = evaluate_observation_term(cf, state)
        
        # For unit departures and unit R: J_o = 0.5 * n_obs
        expected_cost = 0.5 * n_obs
        @test obs_cost ≈ expected_cost atol=1e-10
        
        # Test scaling property
        scale_factor = 1.5
        state[1:n_obs] = observations .+ scale_factor
        scaled_cost = evaluate_observation_term(cf, state)
        expected_scaled = 0.5 * n_obs * scale_factor^2
        @test scaled_cost ≈ expected_scaled atol=1e-10
    end
    
    @testset "Total Cost Function Evaluation" begin
        config = AnalysisConfig{Float64}(grid_size = (4, 4, 2))
        cf = CostFunction(config)
        
        n_state = prod(config.grid_size)
        n_obs = 6
        
        # Set up background covariance and background state
        B = Diagonal(ones(n_state))
        initialize_background_covariance!(cf, B)
        background = zeros(n_state)
        set_background!(cf, background)
        
        # Set up observations
        H(x) = x[1:n_obs]
        observations = randn(n_obs)
        R = Diagonal(ones(n_obs))
        
        set_observations!(cf, observations, H)
        cf.observation_covariance = R
        
        # Evaluate total cost
        state = ones(n_state)
        state[1:n_obs] = observations .+ 0.5  # Half unit departure from observations
        
        total_cost = evaluate_cost(cf, state)
        
        # Manual calculation
        bg_cost_manual = 0.5 * sum(state.^2)  # Background departure
        obs_cost_manual = 0.5 * sum(0.5^2 for _ in 1:n_obs)  # Observation departure
        expected_total = bg_cost_manual + obs_cost_manual
        
        @test total_cost ≈ expected_total atol=1e-10
        @test cf.current_cost ≈ total_cost
    end
    
    @testset "Gradient Computation" begin
        config = AnalysisConfig{Float64}(grid_size = (4, 3, 2))
        cf = CostFunction(config)
        
        n_state = prod(config.grid_size)
        n_obs = 4
        
        # Set up problem
        B = Diagonal(ones(n_state))
        initialize_background_covariance!(cf, B)
        background = zeros(n_state)
        set_background!(cf, background)
        
        H(x) = x[1:n_obs]
        observations = zeros(n_obs)
        R = Diagonal(ones(n_obs))
        
        set_observations!(cf, observations, H)
        cf.observation_covariance = R
        
        # Test gradient at a specific point
        state = ones(n_state)
        gradient = compute_gradient(cf, state)
        
        @test length(gradient) == n_state
        @test isfinite(norm(gradient))
        @test cf.current_gradient_norm ≈ norm(gradient)
        
        # For this simple case:
        # Background gradient: B^(-1) * (state - background) = state (since B = I, bg = 0)
        # Observation gradient: H^T * R^(-1) * (H(state) - obs) = [1,1,1,1,0,0,...]
        expected_bg_grad = state
        expected_obs_grad = zeros(n_state)
        expected_obs_grad[1:n_obs] .= 1.0  # H(state) - obs = [1,1,1,1] - [0,0,0,0]
        
        expected_total_grad = expected_bg_grad + expected_obs_grad
        @test gradient ≈ expected_total_grad atol=1e-10
    end
    
    @testset "Hybrid Gradient Computation" begin
        config = AnalysisConfig{Float64}(
            grid_size = (4, 3, 2),
            ensemble_size = 6,
            hybrid_coeff = 0.5
        )
        hcf = HybridCostFunction(config)
        
        n_state = prod(config.grid_size)
        
        # Set up hybrid cost function
        B_static = Diagonal(ones(n_state))
        initialize_background_covariance!(hcf, B_static)
        
        perturbations = randn(n_state, config.ensemble_size)
        initialize_ensemble!(hcf, perturbations)
        
        background = zeros(n_state)
        set_background!(hcf, background)
        
        # Simple observations
        H(x) = x[1:3]
        observations = zeros(3)
        R = Diagonal(ones(3))
        
        set_observations!(hcf, observations, H)
        hcf.observation_covariance = R
        
        # Compute gradient
        state = ones(n_state) * 0.1  # Small state to avoid numerical issues
        gradient = compute_gradient(hcf, state)
        
        @test length(gradient) == n_state
        @test isfinite(norm(gradient))
        @test hcf.current_gradient_norm ≈ norm(gradient)
    end
    
    @testset "Gradient Check" begin
        # Finite difference gradient check
        config = AnalysisConfig{Float64}(grid_size = (3, 2, 2))
        cf = CostFunction(config)
        
        n_state = prod(config.grid_size)
        
        # Set up simple problem
        B = Diagonal(ones(n_state) * 2.0)  # Background error variance = 2
        initialize_background_covariance!(cf, B)
        background = zeros(n_state)
        set_background!(cf, background)
        
        H(x) = [x[1], x[3]]  # Observe first and third elements
        observations = [1.0, -0.5]
        R = Diagonal([0.5, 1.0])  # Different observation errors
        
        set_observations!(cf, observations, H)
        cf.observation_covariance = R
        
        # Test point
        state = randn(n_state) * 0.1
        
        # Analytical gradient
        analytical_grad = compute_gradient(cf, state)
        
        # Finite difference gradient
        ε = 1e-7
        fd_grad = zeros(n_state)
        
        for i in 1:n_state
            state_plus = copy(state)
            state_minus = copy(state)
            state_plus[i] += ε
            state_minus[i] -= ε
            
            cost_plus = evaluate_cost(cf, state_plus)
            cost_minus = evaluate_cost(cf, state_minus)
            
            fd_grad[i] = (cost_plus - cost_minus) / (2ε)
        end
        
        # Compare gradients
        @test analytical_grad ≈ fd_grad atol=1e-5
    end
    
    @testset "Cost Function Properties" begin
        config = AnalysisConfig{Float64}(grid_size = (4, 4, 2))
        cf = CostFunction(config)
        
        n_state = prod(config.grid_size)
        
        # Set up positive definite background covariance
        B = Diagonal(ones(n_state) + 0.1)  # Ensure positive definiteness
        initialize_background_covariance!(cf, B)
        background = zeros(n_state)
        set_background!(cf, background)
        
        H(x) = x[1:5]
        observations = zeros(5)
        R = Diagonal(ones(5) * 0.5)
        
        set_observations!(cf, observations, H)
        cf.observation_covariance = R
        
        # Test non-negativity
        for _ in 1:10
            state = randn(n_state)
            cost = evaluate_cost(cf, state)
            @test cost >= 0.0
        end
        
        # Test that minimum is at background when observations match H(background)
        cost_at_bg = evaluate_cost(cf, background)
        nearby_state = background + randn(n_state) * 0.01
        cost_nearby = evaluate_cost(cf, nearby_state)
        @test cost_at_bg <= cost_nearby  # Should be minimum (or very close)
    end
    
    @testset "Error Handling" begin
        config = AnalysisConfig{Float64}(grid_size = (4, 4, 2))
        cf = CostFunction(config)
        
        n_state = prod(config.grid_size)
        
        # Test with uninitialized components
        state = randn(n_state)
        
        # Should handle missing background covariance gracefully
        @test_nowarn evaluate_cost(cf, state)
        
        # Test dimension mismatches
        wrong_size_B = Diagonal(ones(n_state + 5))
        @test_throws BoundsError initialize_background_covariance!(cf, wrong_size_B)
    end
    
    @testset "Numerical Stability" begin
        config = AnalysisConfig{Float64}(grid_size = (4, 4, 2))
        cf = CostFunction(config)
        
        n_state = prod(config.grid_size)
        
        # Test with very small covariance values
        B_small = Diagonal(ones(n_state) * 1e-10)
        initialize_background_covariance!(cf, B_small)
        set_background!(cf, zeros(n_state))
        
        H(x) = x[1:3]
        set_observations!(cf, zeros(3), H)
        cf.observation_covariance = Diagonal(ones(3) * 1e-10)
        
        state = ones(n_state) * 1e-5
        cost = evaluate_cost(cf, state)
        gradient = compute_gradient(cf, state)
        
        @test isfinite(cost)
        @test all(isfinite.(gradient))
        
        # Test with large values
        B_large = Diagonal(ones(n_state) * 1e10)
        initialize_background_covariance!(cf, B_large)
        
        large_state = ones(n_state) * 1e3
        cost_large = evaluate_cost(cf, large_state)
        gradient_large = compute_gradient(cf, large_state)
        
        @test isfinite(cost_large)
        @test all(isfinite.(gradient_large))
    end
end

println("CostFunctions module tests completed successfully.")